import ast
import json
import javalang
import esprima
import argparse
import os
import pickle
import time
from pathlib import Path
import torch.multiprocessing as mp
import os.path as osp
import pandas as pd
import networkx as nx

VERSION = "v2.3"
NODE_TYPE_DIRECTORY = "directory"
NODE_TYPE_FILE = "file"
NODE_TYPE_CLASS = "class"
NODE_TYPE_FUNCTION = "function"

EDGE_TYPE_CONTAINS = "contains"
EDGE_TYPE_INHERITS = "inherits"
EDGE_TYPE_INVOKES = "invokes"
EDGE_TYPE_IMPORTS = "imports"


def get_inner_nodes(query_node, src_node, graph):
    """获取源节点内部的所有节点"""
    inner_nodes = []
    for _, dst_node, attr in graph.edges(src_node, data=True):
        if attr.get("type") == EDGE_TYPE_CONTAINS and dst_node != query_node:
            inner_nodes.append(dst_node)
            if graph.nodes[dst_node].get("type") == NODE_TYPE_CLASS:
                inner_nodes.extend(get_inner_nodes(query_node, dst_node, graph))
    return inner_nodes


def find_all_possible_callee(node, graph):
    """查找节点可能调用的所有函数和类"""
    callee_nodes, callee_alias = [], {}
    cur_node = node
    pre_node = node

    def find_parent(_cur_node):
        for predecessor in graph.predecessors(_cur_node):
            edges_data = graph.get_edge_data(predecessor, _cur_node)
            if edges_data:
                for key, attr in edges_data.items():
                    if attr.get("type") == EDGE_TYPE_CONTAINS:
                        return predecessor

    while True:
        callee_nodes.extend(get_inner_nodes(pre_node, cur_node, graph))

        if graph.nodes[cur_node].get("type") == NODE_TYPE_FILE:
            # 检查递归导入的文件
            file_list = []
            file_stack = [cur_node]
            while file_stack:
                for _, dst_node, attr in graph.edges(file_stack.pop(), data=True):
                    if attr.get(
                        "type"
                    ) == EDGE_TYPE_IMPORTS and dst_node not in file_list + [cur_node]:
                        if graph.nodes[dst_node].get(
                            "type"
                        ) == NODE_TYPE_FILE and dst_node.endswith("__init__.py"):
                            file_list.append(dst_node)
                            file_stack.append(dst_node)

            for file in file_list:
                callee_nodes.extend(get_inner_nodes(cur_node, file, graph))
                for _, dst_node, attr in graph.edges(file, data=True):
                    if attr.get("type") == EDGE_TYPE_IMPORTS:
                        if attr.get("alias") is not None:
                            callee_alias[attr["alias"]] = dst_node
                        if graph.nodes[dst_node].get("type") in [
                            NODE_TYPE_FILE,
                            NODE_TYPE_CLASS,
                        ]:
                            callee_nodes.extend(get_inner_nodes(file, dst_node, graph))
                        if graph.nodes[dst_node].get("type") in [
                            NODE_TYPE_FUNCTION,
                            NODE_TYPE_CLASS,
                        ]:
                            callee_nodes.append(dst_node)

            # 检查导入的函数和类
            for _, dst_node, attr in graph.edges(cur_node, data=True):
                if attr.get("type") == EDGE_TYPE_IMPORTS:
                    if attr.get("alias") is not None:
                        callee_alias[attr["alias"]] = dst_node
                    if graph.nodes[dst_node].get("type") in [
                        NODE_TYPE_FILE,
                        NODE_TYPE_CLASS,
                    ]:
                        callee_nodes.extend(get_inner_nodes(cur_node, dst_node, graph))
                    if graph.nodes[dst_node].get("type") in [
                        NODE_TYPE_FUNCTION,
                        NODE_TYPE_CLASS,
                    ]:
                        callee_nodes.append(dst_node)
            break

        pre_node = cur_node
        cur_node = find_parent(cur_node)

    return callee_nodes, callee_alias


class JavaAnalyzer:
    def __init__(self, filename, source_code, project_root):
        self.filename = filename
        self.source_code = source_code
        self.source_lines = source_code.split("\n")
        self.project_root = project_root
        self.nodes = []
        self.edges = []
        self.package_name = ""
        self.imports = []

    def _get_file_path_identifier(self):
        """获取文件路径标识符，格式为：文件夹/文件"""
        rel_path = os.path.relpath(self.filename, self.project_root)
        return rel_path.replace(os.sep, "/").replace(".java", "")

    def _create_node_identifier(self, class_name=None, method_name=None):
        """创建节点标识符，格式为：文件夹/文件:类.函数"""
        file_id = self._get_file_path_identifier()
        if class_name and method_name:
            return f"{file_id}:{class_name}.{method_name}"
        elif class_name:
            return f"{file_id}:{class_name}"
        else:
            return file_id

    def analyze(self, tree):
        """分析Java语法树"""
        if hasattr(tree, "package") and tree.package:
            self.package_name = tree.package.name

        if hasattr(tree, "imports"):
            for import_decl in tree.imports:
                self.imports.append(import_decl.path)

        if hasattr(tree, "types"):
            for type_decl in tree.types:
                self._visit_type(type_decl)

        return {"nodes": self.nodes, "edges": self.edges, "imports": self.imports}

    def _get_line_number(self, node):
        """获取节点行号"""
        try:
            if hasattr(node, "position") and node.position:
                return node.position.line
        except:
            pass
        return 0

    def _extract_code_block(self, start_line, end_line=None):
        """提取代码块"""
        if start_line <= 0 or start_line > len(self.source_lines):
            return ""

        if end_line is None:
            end_line = start_line

        start_idx = start_line - 1
        end_idx = min(end_line, len(self.source_lines))

        return "\n".join(self.source_lines[start_idx:end_idx])

    def _visit_type(self, node):
        """访问类型声明"""
        if isinstance(node, javalang.tree.ClassDeclaration):
            self._visit_class(node)
        elif isinstance(node, javalang.tree.InterfaceDeclaration):
            self._visit_interface(node)

    def _visit_class(self, node):
        """访问类声明"""
        class_name = node.name
        class_identifier = self._create_node_identifier(class_name)

        start_line = self._get_line_number(node)
        code_block = self._extract_code_block(start_line)

        class_node = {
            "name": class_identifier,
            "type": NODE_TYPE_CLASS,
            "code": code_block,
            "start_line": start_line,
            "file_path": self.filename,
        }

        self.nodes.append(class_node)

        # 处理类成员
        if hasattr(node, "body"):
            for member in node.body:
                if isinstance(member, javalang.tree.MethodDeclaration):
                    self._visit_method(member, class_name)
                elif isinstance(member, javalang.tree.ConstructorDeclaration):
                    self._visit_constructor(member, class_name)

    def _visit_interface(self, node):
        """访问接口声明"""
        interface_name = node.name
        interface_identifier = self._create_node_identifier(interface_name)

        start_line = self._get_line_number(node)
        code_block = self._extract_code_block(start_line)

        interface_node = {
            "name": interface_identifier,
            "type": NODE_TYPE_CLASS,
            "code": code_block,
            "start_line": start_line,
            "file_path": self.filename,
            "is_interface": True,
        }

        self.nodes.append(interface_node)

        if hasattr(node, "body"):
            for member in node.body:
                if isinstance(member, javalang.tree.MethodDeclaration):
                    self._visit_method(member, interface_name)

    def _visit_method(self, node, parent_class):
        """访问方法声明"""
        method_name = node.name
        method_identifier = self._create_node_identifier(parent_class, method_name)
        class_identifier = self._create_node_identifier(parent_class)

        start_line = self._get_line_number(node)
        code_block = self._extract_code_block(start_line)

        parameters = []
        if hasattr(node, "parameters") and node.parameters:
            for param in node.parameters:
                parameters.append(
                    {
                        "name": param.name,
                        "type": str(param.type) if param.type else "unknown",
                    }
                )

        method_node = {
            "name": method_identifier,
            "type": NODE_TYPE_FUNCTION,
            "code": code_block,
            "start_line": start_line,
            "file_path": self.filename,
            "parameters": parameters,
            "return_type": str(node.return_type) if node.return_type else "void",
        }

        self.nodes.append(method_node)

        # 添加包含关系
        self.edges.append(
            {
                "source": class_identifier,
                "target": method_identifier,
                "type": EDGE_TYPE_CONTAINS,
            }
        )

    def _visit_constructor(self, node, parent_class):
        """访问构造函数"""
        constructor_name = node.name
        constructor_identifier = self._create_node_identifier(
            parent_class, constructor_name
        )
        class_identifier = self._create_node_identifier(parent_class)

        start_line = self._get_line_number(node)
        code_block = self._extract_code_block(start_line)

        parameters = []
        if hasattr(node, "parameters") and node.parameters:
            for param in node.parameters:
                parameters.append(
                    {
                        "name": param.name,
                        "type": str(param.type) if param.type else "unknown",
                    }
                )

        constructor_node = {
            "name": constructor_identifier,
            "type": NODE_TYPE_FUNCTION,
            "code": code_block,
            "start_line": start_line,
            "file_path": self.filename,
            "parameters": parameters,
            "is_constructor": True,
        }

        self.nodes.append(constructor_node)

        # 添加包含关系
        self.edges.append(
            {
                "source": class_identifier,
                "target": constructor_identifier,
                "type": EDGE_TYPE_CONTAINS,
            }
        )


def analyze_file_java(filepath, project_root):
    """分析Java文件"""
    try:
        with open(filepath, "r", encoding="utf-8") as file:
            code = file.read()

        tree = javalang.parse.parse(code)
        analyzer = JavaAnalyzer(filepath, code, project_root)
        return analyzer.analyze(tree)
    except Exception as e:
        print(f"Error parsing {filepath}: {e}")
        return {"nodes": [], "edges": [], "imports": []}


def build_graph_from_data(nodes, edges):
    """从节点和边数据构建NetworkX图"""
    graph = nx.DiGraph()

    for node in nodes:
        graph.add_node(node["name"], **node)

    for edge in edges:
        graph.add_edge(edge["source"], edge["target"], **edge)

    return graph


def analyze_java_project(project_path):
    """分析Java项目"""
    all_nodes = []
    all_edges = []
    project_path = os.path.abspath(project_path)

    # 遍历项目目录
    for root, dirs, files in os.walk(project_path):
        relative_dir_path = os.path.relpath(root, project_path).replace(os.sep, "/")

        # 添加目录节点
        if relative_dir_path != ".":
            dir_node = {
                "name": relative_dir_path,
                "type": NODE_TYPE_DIRECTORY,
                "file_path": relative_dir_path,
            }
            all_nodes.append(dir_node)

            # 链接到父目录
            parent_dir_path = os.path.dirname(relative_dir_path)
            if parent_dir_path and parent_dir_path != ".":
                all_edges.append(
                    {
                        "source": parent_dir_path,
                        "target": relative_dir_path,
                        "type": EDGE_TYPE_CONTAINS,
                    }
                )

        # 处理文件
        for file_name in files:
            if file_name.endswith(".java"):
                file_path = os.path.join(root, file_name)
                relative_file_path = os.path.relpath(file_path, project_path).replace(
                    os.sep, "/"
                )

                try:
                    with open(file_path, "r", encoding="utf-8") as f:
                        file_content = f.read()
                except Exception:
                    file_content = ""

                # 添加文件节点
                file_node = {
                    "name": relative_file_path.replace(".java", ""),
                    "type": NODE_TYPE_FILE,
                    "code": file_content,
                    "file_path": relative_file_path,
                }
                all_nodes.append(file_node)

                # 链接目录到文件
                if relative_dir_path != ".":
                    all_edges.append(
                        {
                            "source": relative_dir_path,
                            "target": relative_file_path.replace(".java", ""),
                            "type": EDGE_TYPE_CONTAINS,
                        }
                    )

                # 分析Java文件内容
                analysis_result = analyze_file_java(file_path, project_path)

                all_nodes.extend(analysis_result["nodes"])
                all_edges.extend(analysis_result["edges"])

                # 将文件链接到其顶层代码元素
                for node in analysis_result["nodes"]:
                    if node.get("type") == NODE_TYPE_CLASS:
                        all_edges.append(
                            {
                                "source": relative_file_path.replace(".java", ""),
                                "target": node["name"],
                                "type": EDGE_TYPE_CONTAINS,
                            }
                        )

    return {"nodes": all_nodes, "edges": all_edges}


def find_callees_for_function(function_name, project_data):
    """为指定函数查找所有可能的被调用者"""
    graph = build_graph_from_data(project_data["nodes"], project_data["edges"])

    # 查找匹配的函数节点
    function_node = None
    for node_name, node_data in graph.nodes(data=True):
        if node_data.get("type") == NODE_TYPE_FUNCTION and function_name in node_name:
            function_node = node_name
            break

    if function_node:
        callee_nodes, callee_alias = find_all_possible_callee(function_node, graph)
        return {
            "function": function_node,
            "callees": callee_nodes,
            "aliases": callee_alias,
        }

    return {"function": None, "callees": [], "aliases": {}}


def save_to_excel(data, output_path):
    """保存数据到Excel文件"""
    df = pd.DataFrame(data)
    if not df.empty:
        # 处理复杂字段
        for col in df.columns:
            df[col] = df[col].apply(
                lambda x: (
                    json.dumps(x, ensure_ascii=False)
                    if isinstance(x, (list, dict))
                    else x
                )
            )
        df.to_excel(output_path, index=False, engine="openpyxl")
        print(f"Saved {len(df)} records to {output_path}")


def run(rank, repo_name, repo_path, out_path, download_repo=False, instance_data=None):
    """处理单个仓库"""
    output_file = f"{osp.join(out_path, repo_name)}.pkl"
    repo_dir = osp.join(repo_path, repo_name)

    print(f"[{rank}] Start process {repo_name}")
    try:
        if os.path.exists(repo_dir):
            # 分析项目
            result = analyze_java_project(repo_dir)

            # 保存原始数据
            with open(output_file, "wb") as f:
                pickle.dump(result, f)

            # 保存节点和边信息
            if result["nodes"]:
                save_to_excel(
                    result["nodes"], f"{osp.join(out_path, repo_name)}_nodes.xlsx"
                )
            if result["edges"]:
                save_to_excel(
                    result["edges"], f"{osp.join(out_path, repo_name)}_edges.xlsx"
                )

            # 示例：为项目中的第一个函数查找可能的被调用者
            functions = [
                node
                for node in result["nodes"]
                if node.get("type") == NODE_TYPE_FUNCTION
            ]
            if functions:
                first_function = functions[0]["name"]
                callees_info = find_callees_for_function(first_function, result)
                print(
                    f"[{rank}] Function {first_function} can call: {len(callees_info['callees'])} items"
                )

            print(
                f"[{rank}] Processed {repo_name} - Found {len(result['nodes'])} nodes and {len(result['edges'])} edges"
            )
        else:
            print(f"[{rank}] Repository directory not found: {repo_dir}")
    except Exception as e:
        print(f"[{rank}] Error processing {repo_name}: {e}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--repo_path", type=str, default="/data_ext/ref_code/LocAgent/data/data/pp_data"
    )
    parser.add_argument("--index_dir", type=str, default="index_data")
    args = parser.parse_args()

    dataset_name = "pp_data"
    args.index_dir = f"{args.index_dir}/{dataset_name}/graph_index_{VERSION}/"
    os.makedirs(args.index_dir, exist_ok=True)

    start_time = time.time()

    # 测试单个仓库
    repo_name = "pingp-auto-test"
    run(0, repo_name, args.repo_path, args.index_dir)

    end_time = time.time()
    print(f"Total Execution time = {end_time - start_time:.3f}s")
